#pragma OPENCL EXTENSION cl_khr_3d_image_writes : enable
#pragma OPENCL EXTENSION cl_arm_printf : enable
__constant sampler_t smp_none = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;

__kernel void conv_1x1_image2d(__write_only image2d_t output_tex2d,                     // 32个36个14400行
                               __read_only image2d_t input_tex2d,
                               __global float16 *weight_buffer,                         // 32*36*32*2*4*4
                               int4 input_sliced_shape, int4 weight_sliced_shape) {     // 1, 36, 14400, 32;   64, 6, 6, 32
    int x_idx = get_global_id(0) * 2;       // 0~14400,         dst_width
    int y_idx = get_global_id(1);           // 0~36,            dst_height
    int slice_idx = get_global_id(2) * 2;   // grp_ids: 0~64    dst_ocs

    int src_width = input_sliced_shape.z, src_height = input_sliced_shape.y, src_slice = input_sliced_shape.w;
    int weight_o_slice = weight_sliced_shape.x, weight_i_slice = weight_sliced_shape.w;
    int weight_height = weight_sliced_shape.y, weight_width = weight_sliced_shape.z; 
    int weight_kernel_size = weight_height * weight_width;
    int dst_width = src_width, dst_height = src_height, dst_slice = weight_o_slice;

    if (x_idx >= dst_width || y_idx >= dst_height || slice_idx >= dst_slice) return;

    __global float16 *weight_start = weight_buffer + ((slice_idx/2)*weight_kernel_size+y_idx)*weight_i_slice*2;
    float4 v0 = 0, v1 = 0, group1_v0 = 0, group1_v1 = 0;
    int global_y_idx = y_idx;
    for (int ic=0; ic<src_slice; ++ic) {
        float4 input0 = convert_float4(read_imagef(input_tex2d, smp_none, (int2)(x_idx, global_y_idx)));
        float4 input1 = convert_float4(read_imagef(input_tex2d, smp_none, (int2)(x_idx+1, global_y_idx)));
        // float4 input0 = input_buffer[input_idx], input1 = input_buffer[input_idx + 1];
        float16 weight0 = weight_start[0], weight1 = weight_start[1];
        v0 += weight0.s0123 * input0.s0;                           // weight0.s0123是4个filter中的channel_s同一个位置, s0.s0是input中channel_s上同一个位置的值，这样子算出oc4
        v0 += weight0.s4567 * input0.s1;                           // weight0.s4567是4个filter中的channel_s+1的同一个位置，这里是沿着ic维度的累加
        v0 += weight0.s89ab * input0.s2;
        v0 += weight0.scdef * input0.s3;                           // 直到channel_s+3，到此一个ic4计算完成
        v1 += weight0.s0123 * input1.s0;                           // 算input中的下一个ic4
        v1 += weight0.s4567 * input1.s1;
        v1 += weight0.s89ab * input1.s2;
        v1 += weight0.scdef * input1.s3;

        group1_v0 += weight1.s0123 * input0.s0;
        group1_v0 += weight1.s4567 * input0.s1;
        group1_v0 += weight1.s89ab * input0.s2;
        group1_v0 += weight1.scdef * input0.s3;
        group1_v1 += weight1.s0123 * input1.s0;
        group1_v1 += weight1.s4567 * input1.s1;
        group1_v1 += weight1.s89ab * input1.s2;
        group1_v1 += weight1.scdef * input1.s3;

        global_y_idx = global_y_idx + src_height;
        weight_start = weight_start + 2;
    }

    global_y_idx = slice_idx * dst_height + y_idx;
    write_imagef(output_tex2d, (int2)(x_idx, global_y_idx), v0);
    write_imagef(output_tex2d, (int2)(x_idx + 1, global_y_idx), v1);
    
    global_y_idx = global_y_idx + dst_height;
    write_imagef(output_tex2d, (int2)(x_idx, global_y_idx), group1_v0);
    write_imagef(output_tex2d, (int2)(x_idx + 1, global_y_idx), group1_v1);
}

// #pragma OPENCL EXTENSION cl_khr_3d_image_writes : enable
// #pragma OPENCL EXTENSION cl_arm_printf : enable
// __constant sampler_t smp_none = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;

// __kernel void conv_1x1_image2d(__write_only image2d_t output_tex2d,                     // 32个36个14400行
//                                __read_only image2d_t input_tex2d,
//                                __global float16 *weight_buffer,                         // 32*36*32*2*4*4
//                                int4 input_sliced_shape, int4 weight_sliced_shape) {     // 1, 36, 14400, 32;   64, 6, 6, 32
//     int x_idx = get_global_id(0) * 2;       // 0~14400,         dst_width
//     int y_idx = get_global_id(1);           // 0~36,            dst_height
//     int slice_idx = get_global_id(2);       // grp_ids: 0~16    dst_ocs

//     int src_width = input_sliced_shape.z, src_height = input_sliced_shape.y, src_slice = input_sliced_shape.w;
//     int weight_o_slice = weight_sliced_shape.x, weight_i_slice = weight_sliced_shape.w;
//     int weight_height = weight_sliced_shape.y, weight_width = weight_sliced_shape.z; 
//     int weight_kernel_size = weight_height * weight_width;
//     int dst_width = src_width, dst_height = src_height, dst_slice = weight_o_slice;

//     if (x_idx >= dst_width || y_idx >= dst_height || slice_idx >= dst_slice) return;

//     __global float16 *weight_start = weight_buffer + (slice_idx*weight_kernel_size+y_idx)*weight_i_slice;
//     float4 v0 = 0, v1 = 0;
//     int global_y_idx = y_idx;
//     for (int ic=0; ic<src_slice; ++ic) {
//         float4 input0 = convert_float4(read_imagef(input_tex2d, smp_none, (int2)(x_idx, global_y_idx)));
//         float4 input1 = convert_float4(read_imagef(input_tex2d, smp_none, (int2)(x_idx+1, global_y_idx)));
//         // float4 input0 = input_buffer[input_idx], input1 = input_buffer[input_idx + 1];
//         float16 weight0 = weight_start[0], weight1 = weight_start[1];
//         v0 += weight0.s0123 * input0.s0;                           // weight0.s0123是4个filter中的channel_s同一个位置, s0.s0是input中channel_s上同一个位置的值，这样子算出oc4
//         v0 += weight0.s4567 * input0.s1;                           // weight0.s4567是4个filter中的channel_s+1的同一个位置，这里是沿着ic维度的累加
//         v0 += weight0.s89ab * input0.s2;
//         v0 += weight0.scdef * input0.s3;                           // 直到channel_s+3，到此一个ic4计算完成
//         v1 += weight0.s0123 * input1.s0;                           // 算input中的下一个ic4
//         v1 += weight0.s4567 * input1.s1;
//         v1 += weight0.s89ab * input1.s2;
//         v1 += weight0.scdef * input1.s3;

//         global_y_idx = global_y_idx + src_height;
//         weight_start = weight_start + 1;
//     }

//     global_y_idx = slice_idx * dst_height + y_idx;
//     write_imagef(output_tex2d, (int2)(x_idx, global_y_idx), v0);
//     write_imagef(output_tex2d, (int2)(x_idx + 1, global_y_idx), v1);
// }